//
// @Auther: by wujehy 
// @Email:wujehy@qq.com
// @Data:20-4-18
// @Time:下午1:38
//

#include "SecurityCommon.h"
#include <openssl/ossl_typ.h>
#include <openssl/evp.h>
#include <openssl/hmac.h>
#include <openssl/rand.h>



#define HASH_OUTPUT_SIZE 32

#define MIN(a, b) (((a)<(b))?(a):(b))
#define MAX(a, b) (((a)>(b))?(a):(b))

using namespace GeeJoan::Common;

int GeeJoan::Common::common_random_generator(uint8_t *data, size_t len, void *user_data)
{
    if (RAND_bytes( data, len ))
    {
        return 0;
    } else
    {
        return -1;
    }
}

int64_t GeeJoan::Common::jenkins_hash(const char *key, size_t len)
{
    uint64_t hash, i;
    for (hash = i = 0; i < len; ++i)
    {
        hash += key[i];
        hash += (hash << 10);
        hash ^= (hash >> 6);
    }
    hash += (hash << 3);
    hash ^= (hash >> 11);
    hash += (hash << 15);
    return hash;
}

const EVP_CIPHER *aes_cipher(int cipher, size_t key_len)
{
    if (cipher == CIPHER_AES_CBC_PKCS5)
    {
        if (key_len == 16)
        {
            return EVP_aes_128_cbc();
        } else if (key_len == 24)
        {
            return EVP_aes_192_cbc();
        } else if (key_len == 32)
        {
            return EVP_aes_256_cbc();
        }
    } else if (cipher == CIPHER_AES_CTR_NOPADDING)
    {
        if (key_len == 16)
        {
            return EVP_aes_128_ctr();
        } else if (key_len == 24)
        {
            return EVP_aes_192_ctr();
        } else if (key_len == 32)
        {
            return EVP_aes_256_ctr();
        }
    }
    return 0;
}

int GeeJoan::Common::common_encrypt(
        std::string &outBuffer, CIPHER_AES_TYPE cipher, const uint8_t *key, size_t key_len, const uint8_t *iv, size_t iv_len,
        const uint8_t *plaintext, size_t plaintext_len, void *user_data)
{
    int result = 0;
    bool isComplete = false;
    EVP_CIPHER_CTX *ctx = 0;
    uint8_t *out_buf = 0;

    const EVP_CIPHER *evp_cipher = aes_cipher( cipher, key_len );
    if (!evp_cipher)
    {
        fprintf( stderr, "invalid AES mode or key size: %zu\n", key_len );
        return COMMON_ERROR_CODE_UNKNOW;
    }

    if (iv_len != 16)
    {
        fprintf( stderr, "invalid AES IV size: %zu\n", iv_len );
        return COMMON_ERROR_CODE_UNKNOW;
    }

    if (plaintext_len > INT_MAX - EVP_CIPHER_block_size( evp_cipher ))
    {
        fprintf( stderr, "invalid plaintext length: %zu\n", plaintext_len );
        return COMMON_ERROR_CODE_UNKNOW;
    }

#if OPENSSL_VERSION_NUMBER >= 0x1010000fL
    if (!isComplete)
    {
        ctx = EVP_CIPHER_CTX_new();
        if (!ctx)
        {
            result = COMMON_ERROR_CODE_NOMEM;
            isComplete = true;
        }
    }

#else
    if(!isComplete){
        ctx = malloc(sizeof(EVP_CIPHER_CTX));
        if(!ctx) {
            result = COMMON_ERROR_CODE_NOMEM;
            isComplete=true;
        }
        EVP_CIPHER_CTX_init(ctx);
    }

#endif


    if (!isComplete)
    {
        result = EVP_EncryptInit_ex( ctx, evp_cipher, 0, key, iv );
        if (!result)
        {
            fprintf( stderr, "cannot initialize cipher\n" );
            result = COMMON_ERROR_CODE_UNKNOW;
            isComplete = true;
        }
    }

    if (!isComplete)
    {
        if (cipher == CIPHER_AES_CTR_NOPADDING)
        {
            result = EVP_CIPHER_CTX_set_padding( ctx, 0 );
            if (!result)
            {
                fprintf( stderr, "cannot set padding\n" );
                result = COMMON_ERROR_CODE_UNKNOW;
                isComplete = true;
            }
        }
    }


    if (!isComplete)
    {
        out_buf = static_cast<uint8_t *>(malloc(
                sizeof( uint8_t ) * (plaintext_len + EVP_CIPHER_block_size( evp_cipher ))));
        if (!out_buf)
        {
            fprintf( stderr, "cannot allocate output buffer\n" );
            result = COMMON_ERROR_CODE_NOMEM;
            isComplete = true;
        }

    }

    int out_len = 0;
    if (!isComplete)
    {

        result = EVP_EncryptUpdate( ctx, out_buf, &out_len, plaintext, plaintext_len );
        if (!result)
        {
            fprintf( stderr, "cannot encrypt plaintext\n" );
            result = COMMON_ERROR_CODE_UNKNOW;
            isComplete = true;
        }
    }

    int final_len = 0;
    if (!isComplete)
    {
        result = EVP_EncryptFinal_ex( ctx, out_buf + out_len, &final_len );
        if (!result)
        {
            fprintf( stderr, "cannot finish encrypting plaintext\n" );
            result = COMMON_ERROR_CODE_UNKNOW;
            isComplete = true;
        }
    }

    if (!isComplete)
    {
        outBuffer = std::string((char *) out_buf, out_len + final_len );
    }

    if (ctx)
    {
#if OPENSSL_VERSION_NUMBER >= 0x1010000fL
        EVP_CIPHER_CTX_free( ctx );
#else
        EVP_CIPHER_CTX_cleanup(ctx);
        free(ctx);
#endif
    }
    if (out_buf)
    {
        free( out_buf );
    }
    return result;
}

int GeeJoan::Common::common_decrypt(
        std::string &outBuffer, CIPHER_AES_TYPE cipher, const uint8_t *key, size_t key_len, const uint8_t *iv, size_t iv_len,
        const uint8_t *ciphertext, size_t ciphertext_len, void *user_data)
{
    int result = 0;
    EVP_CIPHER_CTX *ctx = 0;
    uint8_t *out_buf = 0;
    bool isComplete = false;
    const EVP_CIPHER *evp_cipher = aes_cipher(cipher, key_len);
    if (!evp_cipher)
    {
        fprintf(stderr, "invalid AES mode or key size: %zu\n", key_len);
        return COMMON_ERROR_CODE_INVAL;
    }

    if (iv_len != 16)
    {
        fprintf(stderr, "invalid AES IV size: %zu\n", iv_len);
        return COMMON_ERROR_CODE_INVAL;
    }

    if (ciphertext_len > INT_MAX - EVP_CIPHER_block_size(evp_cipher))
    {
        fprintf(stderr, "invalid ciphertext length: %zu\n", ciphertext_len);
        return COMMON_ERROR_CODE_UNKNOW;
    }
    if (!isComplete)
    {
#if OPENSSL_VERSION_NUMBER >= 0x1010000fL
        ctx = EVP_CIPHER_CTX_new();
        if (!ctx)
        {
            result = COMMON_ERROR_CODE_NOMEM;
            isComplete = true;
        }
#else
        ctx = malloc(sizeof(EVP_CIPHER_CTX));
        if(!ctx) {
            result = COMMON_ERROR_CODE_NOMEM;
            isComplete=true;
        }
        EVP_CIPHER_CTX_init(ctx);
#endif
    }

    if (!isComplete)
    {
        result = EVP_DecryptInit_ex(ctx, evp_cipher, 0, key, iv);
        if (!result)
        {
            fprintf(stderr, "cannot initialize cipher\n");
            result = COMMON_ERROR_CODE_UNKNOW;
            isComplete = true;
        }
    }

    if (!isComplete)
    {
        if (cipher == CIPHER_AES_CTR_NOPADDING)
        {
            result = EVP_CIPHER_CTX_set_padding(ctx, 0);
            if (!result)
            {
                fprintf(stderr, "cannot set padding\n");
                result = COMMON_ERROR_CODE_UNKNOW;
                isComplete = true;
            }
        }
    }

    if (!isComplete)
    {
        out_buf = static_cast<uint8_t *>(malloc(
                sizeof(uint8_t) * (ciphertext_len + EVP_CIPHER_block_size(evp_cipher))));
        if (!out_buf)
        {
            fprintf(stderr, "cannot allocate output buffer\n");
            result = COMMON_ERROR_CODE_UNKNOW;

            isComplete = true;
        }

    }
    int out_len = 0;
    if (!isComplete)
    {
        result = EVP_DecryptUpdate(ctx,
                                   out_buf, &out_len, ciphertext, ciphertext_len);
        if (!result)
        {
            fprintf(stderr, "cannot decrypt ciphertext\n");
            result = COMMON_ERROR_CODE_UNKNOW;
            isComplete = true;
        }
    }


    int final_len = 0;
    if (!isComplete)
    {
        result = EVP_DecryptFinal_ex(ctx, out_buf + out_len, &final_len);
        if (!result)
        {
            fprintf(stderr, "cannot finish decrypting ciphertext\n");
            result = COMMON_ERROR_CODE_UNKNOW;
            isComplete = true;
        }

    }
    if (!isComplete)
    {
//        *output = signal_buffer_create(out_buf, out_len + final_len);
        outBuffer = std::string((char*)out_buf, out_len + final_len);
    }

    if (ctx)
    {
#if OPENSSL_VERSION_NUMBER >= 0x1010000fL
        EVP_CIPHER_CTX_free(ctx);
#else
        EVP_CIPHER_CTX_cleanup(ctx);
        free(ctx);
#endif
    }
    if (out_buf)
    {
        free(out_buf);
    }
    return result;
}

// sha 256 的封装


int GeeJoan::Common::common_hmac_sha256_init(void **hmac_context, const std::string &key, void *user_data)
{

#if OPENSSL_VERSION_NUMBER >= 0x1010000fL
    HMAC_CTX *ctx = HMAC_CTX_new();
    if (!ctx)
    {
        return COMMON_ERROR_CODE_NOMEM;
    }
#else
    HMAC_CTX *ctx = malloc(sizeof(HMAC_CTX));
    if(!ctx) {
        return COMMON_ERROR_CODE_NOMEM;
    }
    HMAC_CTX_init(ctx);
#endif

    *hmac_context = ctx;

    if (HMAC_Init_ex(ctx, (void *) key.c_str(), key.size(), EVP_sha256(), 0) != 1)
    {
        return COMMON_ERROR_CODE_UNKNOW;
    }

    return COMMON_ERROR_CODE_SUCCESS;

}

int GeeJoan::Common::common_hmac_sha256_update(void *hmac_context, const std::string &data, void *user_data)
{
    HMAC_CTX *ctx = (HMAC_CTX *) hmac_context;
    int result = HMAC_Update(ctx, (uint8_t *) data.c_str(), data.size());
    return (result == 1) ? COMMON_ERROR_CODE_SUCCESS : COMMON_ERROR_CODE_FAIL;
}

int GeeJoan::Common::common_hmac_sha256_final(void *hmac_context, std::string &outBuffer, void *user_data)
{
    int result = COMMON_ERROR_CODE_SUCCESS;

    bool isComplete = false; // 标记是否结束
    unsigned char md[EVP_MAX_MD_SIZE];
    unsigned int len = 0;
    HMAC_CTX *ctx = (HMAC_CTX *) hmac_context;

    if (HMAC_Final(ctx, md, &len) != 1)
    {
        return COMMON_ERROR_CODE_UNKNOW;
    }

    outBuffer = std::string((char *) md, len);
    if (outBuffer.empty())
    {
        result = COMMON_ERROR_CODE_NOMEM;
    }

    return result;
}

void GeeJoan::Common::common_hmac_sha256_cleanup(void *hmac_context, void *user_data)
{
    if (hmac_context)
    {
        HMAC_CTX *ctx = (HMAC_CTX *) hmac_context;
#if OPENSSL_VERSION_NUMBER >= 0x1010000fL
        HMAC_CTX_free(ctx);
#else
        HMAC_CTX_cleanup(ctx);
        free(ctx);
#endif
    }
}


int GeeJoan::Common::common_hmac_sha256_base(const std::string &key, const std::string &seed, std::string &outBuffer)
{
    void *hmac_context = 0;
    int result = COMMON_ERROR_CODE_SUCCESS;
    switch (1)
    {
        case 1:
        {
            result = common_hmac_sha256_init(&hmac_context, key, nullptr);
            if (result < 0)
            {
                break;
            }
        }

        case 2:
        {
            result = common_hmac_sha256_update(hmac_context, seed, nullptr);
            if (result < 0)
            {
                break;
            }
        }

        case 3:
        {
            result = common_hmac_sha256_final(hmac_context, outBuffer, nullptr);
            if (result < 0)
            {
                break;
            }
        }
    }

    common_hmac_sha256_cleanup(hmac_context, nullptr);

    return result;
}


// sha 512 的封装

int GeeJoan::Common::common_sha512_digest_init(void **digest_context, void *user_data)
{
    int result = 0;
    bool isComplete = false; // 标记是否结束
    EVP_MD_CTX *ctx;

    ctx = EVP_MD_CTX_create();
    if (!ctx)
    {
        result = COMMON_ERROR_CODE_NOMEM;
        isComplete = true;
    }
    // 没有结束 则 继续操作
    if (!isComplete)
    {
        result = EVP_DigestInit_ex(ctx, EVP_sha512(), 0);
        if (result == 1)
        {
            result = COMMON_ERROR_CODE_SUCCESS;
        } else
        {
            result = COMMON_ERROR_CODE_NOMEM;
        }
    }

    if (result < 0)
    {
        if (ctx)
        {
            EVP_MD_CTX_destroy(ctx);
        }
    } else
    {
        *digest_context = ctx;
    }
    return result;
}

int GeeJoan::Common::common_sha512_digest_update(void *digest_context, const std::string data, void *user_data)
{
    EVP_MD_CTX *ctx = (EVP_MD_CTX *) digest_context;

    int result = EVP_DigestUpdate(ctx, data.c_str(), data.size());

    return (result == 1) ? COMMON_ERROR_CODE_SUCCESS : COMMON_ERROR_CODE_NOMEM;

}

int GeeJoan::Common::common_sha512_digest_final(void *digest_context, std::string &outBuffer, void *user_data)
{
    int result = 0;
    bool isComplete = false;
    unsigned char md[EVP_MAX_MD_SIZE];
    unsigned int len = 0;
    EVP_MD_CTX *ctx = (EVP_MD_CTX *) digest_context;

    switch (1)
    {
        case 1:
        {
            result = EVP_DigestFinal_ex(ctx, md, &len);
            if (result == 1)
            {
                result = COMMON_ERROR_CODE_SUCCESS;
            } else
            {
                result = COMMON_ERROR_CODE_NOMEM;
                break;
            }
        }
        case 2:
        {
            result = EVP_DigestInit_ex(ctx, EVP_sha512(), 0);
            if (result == 1)
            {
                result = COMMON_ERROR_CODE_SUCCESS;
            } else
            {
                result = COMMON_ERROR_CODE_NOMEM;
                break;
            }
        }

        case 3:
        {
            outBuffer = std::string((char *) md, len);
            if (outBuffer.empty())
            {
                result = COMMON_ERROR_CODE_NOMEM;
                break;
            } else
            {
                result = COMMON_ERROR_CODE_SUCCESS;
            }
        }
    }


    return result;
}

void GeeJoan::Common::common_sha512_digest_cleanup(void *digest_context, void *user_data)
{
    EVP_MD_CTX *ctx = (EVP_MD_CTX *) digest_context;
    EVP_MD_CTX_destroy(ctx);
}


int GeeJoan::Common::common_sha512_digest_base(const std::string &input, std::string &outBuffer)
{


    void *digest_context = 0;

    int result = COMMON_ERROR_CODE_SUCCESS;
    switch (1)
    {
        case 1:
        {
            result = common_sha512_digest_init(&digest_context);

            if (result < 0) break;
        }
        case 2:
        {
            result = common_sha512_digest_update(digest_context, input);

            if (result < 0) break;
        }
        case 3:
        {
            result = common_sha512_digest_final(digest_context, outBuffer);
            if (result < 0)break;
        }
    }

    common_sha512_digest_cleanup(digest_context);

    return result;
}
